import gradio as gr
import requests
import json
import base64
from io import BytesIO
from PIL import Image
import traceback

# 本地API配置
API_URL = "http://localhost:8800/v1/images/generations"

# 示例提示词
EXAMPLE_PROMPTS = [
    "一只可爱的熊猫坐在竹子上吃竹子，高清摄影",
    "未来城市的夜景，有飞行汽车和霓虹灯，赛博朋克风格",
    "宁静的山水画，有小桥流水和青山环绕，中国风"
]

def generate_image(prompt, negative_prompt="", width=1024, height=1024, steps=8, seed=0):
    """
    调用本地API生成图片
    
    Args:
        prompt (str): 正向提示词
        negative_prompt (str): 负向提示词
        width (int): 图片宽度
        height (int): 图片高度
        steps (int): 推理步数
        seed (int): 随机种子
        
    Returns:
        tuple: (生成的图片, 状态信息)
    """
    if not prompt or prompt.strip() == "":
        return None, "请输入提示词"
    
    try:
        # 准备请求数据
        payload = {
            "prompt": prompt,
            "negative_prompt": negative_prompt,
            "width": width,
            "height": height,
            "num_inference_steps": steps,
            "seed": seed
        }
        
        # 调用本地API
        print(f"正在调用本地API: {API_URL}")
        response = requests.post(API_URL, json=payload, timeout=1800)  # 30分钟超时
        
        if response.status_code == 200:
            result = response.json()
            
            # 解析返回的base64图片数据
            if "data" in result and len(result["data"]) > 0:
                image_data = result["data"][0]["b64_json"]
                # 解码base64图片数据
                image_bytes = base64.b64decode(image_data)
                image = Image.open(BytesIO(image_bytes))
                
                return image, "图片生成成功"
            else:
                return None, "API返回数据格式错误"
        else:
            return None, f"API调用失败，状态码: {response.status_code}，响应: {response.text}"
            
    except requests.exceptions.ConnectionError:
        return None, "连接API失败，请确保本地API服务已启动"
    except requests.exceptions.Timeout:
        return None, "API调用超时，请稍后重试"
    except Exception as e:
        error_msg = f"生成图片时出错: {str(e)}"
        print(error_msg)
        traceback.print_exc()
        return None, error_msg

# 创建Gradio界面
with gr.Blocks(title="文生图应用") as demo:
    gr.Markdown("# 文生图应用")
    gr.Markdown("输入提示词，调用本地API生成对应的图片")
    
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(
                label="提示词",
                placeholder="请输入图片描述...",
                lines=3
            )
            
            negative_prompt_input = gr.Textbox(
                label="负向提示词",
                placeholder="请输入不希望出现在图片中的内容（可选）...",
                lines=2
            )
            
            with gr.Accordion("高级参数", open=False):
                width_input = gr.Slider(
                    minimum=256, 
                    maximum=2048, 
                    value=1024, 
                    step=64, 
                    label="图片宽度"
                )
                
                height_input = gr.Slider(
                    minimum=256, 
                    maximum=2048, 
                    value=1024, 
                    step=64, 
                    label="图片高度"
                )
                
                steps_input = gr.Slider(
                    minimum=1, 
                    maximum=50, 
                    value=8, 
                    step=1, 
                    label="推理步数"
                )
                
                seed_input = gr.Number(
                    label="随机种子 (0表示随机)",
                    value=0,
                    precision=0
                )
            
            # 添加示例提示词按钮
            gr.Markdown("### 示例提示词")
            for i, example in enumerate(EXAMPLE_PROMPTS, 1):
                btn = gr.Button(f"示例 {i}: {example[:30]}..." if len(example) > 30 else f"示例 {i}: {example}")
                btn.click(
                    fn=lambda ex=example: ex, 
                    outputs=prompt_input
                )
            
            generate_btn = gr.Button("生成图片", variant="primary")
            
        with gr.Column():
            image_output = gr.Image(label="生成的图片")
            status_output = gr.Textbox(label="状态信息", max_lines=3)
    
    # 绑定生成按钮事件
    generate_btn.click(
        fn=generate_image,
        inputs=[
            prompt_input, 
            negative_prompt_input,
            width_input,
            height_input,
            steps_input,
            seed_input
        ],
        outputs=[image_output, status_output]
    )

if __name__ == "__main__":
    print("启动文生图Gradio界面...")
    print(f"请确保本地API服务已在 {API_URL} 启动")
    # 修改为明确支持外部网络访问的配置
    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)
